/*
3d_nms (cpu)c++
fast postprocess(CPU)c++
Written by Zhihua Liu
*/

#include <stdio.h>
#include <math.h>
#include <vector>
#include<iostream>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
//#include "py_decode.h"
#include "src_deploy/cc_decode.cpp"
//#include "src_deploy/cc_decode.h"
namespace py = pybind11;


int decode_head_cpu(py::array_t<float> out_cls,
                    py::array_t<float> out_box,
                    py::array_t<float> out_dir,
                    py::array_t<float> out_pos,
                    py::array_t<float> anchors,
                    py::array_t<float> scores,
                    py::array_t<float> Bbox,
                    py::array_t<int> type,
                    py::array_t<float> ClassThreshold,
                    py::array_t<float> voxels_size,
                    py::array_t<float> pts_ranges,
                    float pos_threshold,
                    int kMaxResultSize,
                    int filled_size_,
                    int ClassNumber,
                    int kRadianNumber,
                    int kInfoNumber,
                    int downsize,
                    bool is_hwc) {
    // Get  ptr from python
    py::buffer_info cls_buf = out_cls.request();
    py::buffer_info box_buf = out_box.request();
    py::buffer_info dir_buf = out_dir.request();
    py::buffer_info pos_buf = out_pos.request();
    py::buffer_info sco_buf = scores.request();
    py::buffer_info Bbo_buf = Bbox.request();
    py::buffer_info lab_buf = type.request();
    py::buffer_info anc_buf = anchors.request();
    py::buffer_info Cla_buf = ClassThreshold.request();
    py::buffer_info vox_buf = voxels_size.request();
    py::buffer_info pts_buf = pts_ranges.request();

    int NetLength = cls_buf.shape[0];
    int NetWidth  = cls_buf.shape[1];

    float* cls_ptr = (float*)cls_buf.ptr;
    float* box_ptr = (float*)box_buf.ptr;
    float* dir_ptr = (float*)dir_buf.ptr;
    float* pos_ptr = (float*)pos_buf.ptr;
    float* sco_ptr = (float*)sco_buf.ptr;
    float* Bbo_ptr = (float*)Bbo_buf.ptr;
    int* lab_ptr = (int*)lab_buf.ptr;
    float* anc_ptr = (float*)anc_buf.ptr;
    float* Cla_ptr = (float*)Cla_buf.ptr;
    float* vox_ptr = (float*)vox_buf.ptr;
    float* pts_ptr = (float*)pts_buf.ptr;

//    std::cout<<cls_ptr[0]<<std::endl;
    // Decode by C++
    if(is_hwc){
        filled_size_ = decode(cls_ptr, box_ptr,  dir_ptr, pos_ptr, anc_ptr, Cla_ptr, pos_threshold,
               sco_ptr, Bbo_ptr, lab_ptr, kMaxResultSize,  filled_size_,
               NetLength, NetWidth, ClassNumber, kRadianNumber, kInfoNumber,
               downsize, vox_ptr, pts_ptr);
    }
    else{
        filled_size_ = decode_chw(cls_ptr, box_ptr,  dir_ptr, pos_ptr, anc_ptr, Cla_ptr, pos_threshold,
               sco_ptr, Bbo_ptr, lab_ptr, kMaxResultSize,  filled_size_,
               NetLength, NetWidth, ClassNumber, kRadianNumber, kInfoNumber,
               downsize, vox_ptr, pts_ptr);
    }
//decode(T const* dpuout_cls,T const* dpuout_box, T const* dpuout_dir, T const* dpuout_pos,
//        float const* anchors_, T const* ClassThreshold,T const& pos_threshold,
//        float *scores,float* Bbox,float* type,int const kMaxResultSize, int filled_size_,
//        int const NetLength,int const NetWidth,int const ClassNumber, int kRadianNumber,int kInfoNumber)
     return filled_size_;
}
int nms_cpu_all(py::array_t<float> scores,
            py::array_t<float> Bbox,
            py::array_t<bool> keep,
            float iou_threshold,
            int filled_size,
            int kInfoNumber) {
    // Get  ptr from python
    py::buffer_info buf5 = scores.request();
    py::buffer_info buf6 = Bbox.request();
    py::buffer_info buf7 = keep.request();

    float* ptr5 = (float*)buf5.ptr;
    float* ptr6 = (float*)buf6.ptr;
    bool* ptr7 = (bool*)buf7.ptr;
    // NMS by C++
    nms_cpu( ptr6, ptr5, ptr7, kInfoNumber,iou_threshold, filled_size);
    return 1;
}

